# coding=UTF-8
from const_ import *
with open("./data_test.txt") as f:
    lines = f.readlines()

id2rawtestdata = {}
for i,line in enumerate(lines):
    if i%3==1:
        id = extract_id(line)
        question=extract_question(lines[i-1])
        id2rawtestdata[id]=question

id2sentence_corpus={}
sentence2id_corpus={}
with open("./corpus.txt") as f:
    lines2 = f.readlines()
    for line in lines2:
        temp_index = line.index(" ")
        doc_index = line[1:temp_index]
        doc_text = line[temp_index+1:]
        id2sentence_corpus[doc_index] = doc_text.replace("\n","")
        sentence2id_corpus[doc_text.replace("\n","")] = doc_index


with open("./relation_test_real.txt",mode="w") as f,open("./top10_candidate.txt") as f2:

    id2right_candidate = {}
    id2left = {}

    lines3 = f2.readlines()
    for i,line in enumerate(lines3):
        if i%15==0:
            id = line.split("=")[1].replace("\n","")
            id2right_candidate[id] = []
            for j in range(i+4,i+14):
                candidate = lines3[j].split("\t")[0]
                new_relation = format_relation(candidate).replace(":"," ")
                id2right_candidate[id].append(sentence2id_corpus[new_relation])

    for i,item in enumerate(lines):
        if i%3==0:
            question = extract_question(item)
            id = extract_id(lines[i+1])
            entity_list = predict_question_entities(question)
            question = update_question(question, entity_list)
            doc_id_str = sentence2id_corpus[question]
            id2left[id] = doc_id_str

    for id in id2left:
        for candidate_doc in id2right_candidate[id]:
            f.write("0 ")
            f.write("T"+id2left[id])
            f.write(" ")
            f.write("T"+candidate_doc)
            f.write("\n")
    print()

